package features.advance.leetcode.array.easy;

import java.util.Arrays;
import java.util.Comparator;
import java.util.PriorityQueue;
import java.util.Random;

/**
 * @author LIN
 * @date 2021-06-01
 */
public class Offer40 {

    public static void main(String[] args) {
        Solution solution = new SolutionQuickSort() {
        };
        int[] arr = {3,2,1};
        int k = 2;
        int[] leastNumbers = solution.getLeastNumbers(arr, k);
        System.out.println(Arrays.toString(leastNumbers));
    }

    static class Solution {
        /**
         * 升序排列，取前面k个元素
         * @param arr
         * @param k
         * @return
         */
        public int[] getLeastNumbers(int[] arr, int k) {
            Arrays.sort(arr);
            int[] res = new int[k];
            for (int i = 0; i < res.length; i++) {
                res[i] = arr[i];
            }
            return res;
        }
    }
    static class SolutionHeap extends Solution{
        @Override
        public int[] getLeastNumbers(int[] arr, int k) {
            // 排除k=0的情况
            int[] res = new int[k];
            if(k == 0){
                return res;
            }
            PriorityQueue<Integer> queue = new PriorityQueue(new Comparator<Integer>() {
                @Override
                public int compare(Integer o1, Integer o2) {
                    return o2 - o1;
                }
            });

            for (int i = 0; i < k; i++) {
                queue.offer(arr[i]);
            }
            for (int i = k; i < arr.length; i++) {
                if(queue.peek() > arr[i]){
                    queue.poll();
                    queue.offer(arr[i]);
                }
            }

            for (int i = 0; i < k; i++) {
                res[i] = queue.poll();
            }
            return res;
        }
    }
    static class SolutionQuickSort extends Solution{
        @Override
        public int[] getLeastNumbers(int[] arr, int k) {
            randomizedSelected(arr, 0, arr.length - 1, k);
            int[] vec = new int[k];
            for (int i = 0; i < k; ++i) {
                vec[i] = arr[i];
            }
            return vec;
        }

        private void randomizedSelected(int[] arr, int l, int r, int k) {
            if (l >= r) {
                return;
            }
            int pos = randomizedPartition(arr, l, r);
            int num = pos - l + 1;
            if (k == num) {
                return;
            } else if (k < num) {
                randomizedSelected(arr, l, pos - 1, k);
            } else {
                randomizedSelected(arr, pos + 1, r, k - num);
            }
        }

        // 基于随机的划分
        private int randomizedPartition(int[] nums, int l, int r) {
            int i = new Random().nextInt(r - l + 1) + l;
            swap(nums, r, i);
            return partition(nums, l, r);
        }

        private int partition(int[] nums, int l, int r) {
            int pivot = nums[r];
            int i = l - 1;
            for (int j = l; j <= r - 1; ++j) {
                if (nums[j] <= pivot) {
                    i = i + 1;
                    swap(nums, i, j);
                }
            }
            swap(nums, i + 1, r);
            return i + 1;
        }

        private void swap(int[] nums, int i, int j) {
            int temp = nums[i];
            nums[i] = nums[j];
            nums[j] = temp;
        }
    }
}
